{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Explain *Anything* Like I'm Five:\n",
    "## A Model for Open Domain Long Form Question Answering\n",
    "\n",
    "---  \n",
    "\n",
    "### Table of Contents  \n",
    "\n",
    "1. [**Introduction**](#intro)  \n",
    "    a. [Preliminaries](#prelims)  \n",
    "    b. [Note on Data and Biases](#reddit_biases)\n",
    "2. [**Task and Data Description**](#task_description)  \n",
    "3. [**Sparse Retrieval:** Making Support Documents with ElasticSearch](#elasticsearch)  \n",
    "4. [**Dense Retrieval:** Making Support Documents with an ELI5-Trained Model](#dense_retrieval)  \n",
    "    b. [Contrastive Training with ELI5 In-Batch Negatives](#dense_train)  \n",
    "    c. [Using the Trained Dense Retriever and Wikipedia Index](#dense_use)  \n",
    "    d. [Retrieval Model Evaluation](#dense_eval)  \n",
    "5. [**Generating Answers with a Sequence-to-Sequence Model**](#generation)   \n",
    "\n",
    "---\n",
    "\n",
    "<img src=\"images/choco_bis.svg\" width=\"900\" align=\"center\"/>  \n",
    "\n",
    "---\n",
    "\n",
    "<center> Screenshot of the <a href=\"https://huggingface.co/qa/\">demo</a> for the question: \"How do people make chocolate?\" </center>\n",
    "\n",
    "---  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='intro'></a>\n",
    "# 1. Introduction\n",
    "\n",
    "Imagine that you are taken with a sudden desire to understand **how the fruit of a tropical tree gets transformed into chocolate bars**, or want to understand **the role of fever in the human body's immune response**: how would you go about finding that information?\n",
    "\n",
    "If your specific question has already been asked and answered clearly and succintly on one of the many question answering platforms available on the Internet (such as [**Quora**](https://www.quora.com/How-is-chocolate-made), [**Reddit**](https://www.reddit.com/user/ex_5_libris/comments/9c8gb1/chocolate_how_chocolate_is_made/), or [**Yahoo Answers**](https://answers.yahoo.com/question/index?qid=20070615082202AArsYN1)), you're in luck: modern search engines will probably take you to that pre-existing answer pretty reliably in a matter of a few clicks.  \n",
    "\n",
    "If no one else has asked the exact question you are interested in, however, the process will be a little more involved. You will likely have to collect relevant information from a variety of sources, figure out how these pieces of knowledge fit together in relation to your query, and synthetize a narrative that answers your initial question.\n",
    "\n",
    "Now, wouldn't it be great if your computer could do all of that for you: **gather** the right sources (*e.g.* paragraphs from relevant Wikipedia pages), **synthetize** the information, and **write up** an easy-to-read, original summary of the relevant points? Such a system isn't quite available yet, at least not one that can provide *reliable* information in its summary. Even though current systems excel at finding an extractive span that answers a factoid question in a given document, they still find open-domain settings where a model needs to find its own sources of information and long answer generation challenging.  \n",
    "\n",
    "Thankfully, a number of recent advances in natural language understanding and generation have made working toward solving this problem much easier! These advances include progress in the pre-training (e.g. [BART](https://arxiv.org/abs/1910.13461), [T5](https://arxiv.org/abs/1910.10683)) and evaluation (e.g. for [factuality](https://arxiv.org/abs/2004.04228)) of sequence-to-sequence models for conditional text generation, new ways to use language understanding models to find information in Wikipedia (e.g. [REALM](https://kentonl.com/pub/gltpc.2020.pdf), [DPR](https://arxiv.org/abs/2004.04906)), and a new training dataset introduced in the paper [ELI5: Long Form Question Answering](https://arxiv.org/abs/1907.09190).\n",
    "\n",
    "The [ELI5](https://arxiv.org/abs/1907.09190) dataset was built by gathering questions that were asked by community members of the [r/explainlikeimfive](https://www.reddit.com/r/explainlikeimfive/) subreddit, along with the answers that were provided by other users. The [rules of the subreddit](https://www.reddit.com/r/explainlikeimfive/wiki/detailed_rules) make this data particularly well suited to training a model for abstractive question answering: the questions need to seek an objective explanation about well established facts, and the answers provided need to be understandable to a layperson without any particular knowledge domain.\n",
    "\n",
    "**In this notebook,** we show how we can take advantage of these recent advances to train a **long form question answering** system which takes in a question, fetches 10 relevant passages from a [Wikipedia snapshot](https://www.aclweb.org/anthology/2020.lrec-1.297/), and writes a multi-sentence answer based on the question and retrieved passages. In particular, training embedding-based retrieval models to gather supporting evidence for open-domain questions is relatively new research area: the last few months have seen some significant progress in cases where direct supervision is available, or with extensive task-specific pretraining. Here, we show how the ELI5 dataset allows us to **train a dense retrieval system without access to either**, making dense retrieval models more accessible. See [this presentation](https://docs.google.com/presentation/d/1A5wJEzFYGdNem7egJ-BTm6EMI3jGNe1lalyChYL54gw) from the [Hugging Face reading group](https://github.com/huggingface/awesome-papers) for a non-exhaustive overview of recent work in the field. \n",
    "\n",
    "Follow along to learn about the steps involved and read some background on the state of the art for some related tasks, or go straight to the:  \n",
    "## [**Live Demo!**](https://huggingface.co/qa/)  \n",
    "(And don't forget to scroll down on the left sidebar to show all of the generation options!)\n",
    "\n",
    "<a id='prelims'></a>\n",
    "### 1.a - Preliminaries  \n",
    "\n",
    "The implementation presented here relies on the [Hugging Face](https://huggingface.co/) [🤗transformers](https://github.com/huggingface/transformers) and [🤗nlp](https://github.com/huggingface/nlp) libraries. Wikipedia indexing relies on [ElasticSearch](https://www.elastic.co/elasticsearch) with its [python bindings](https://github.com/elastic/elasticsearch-py) for the sparse version, and [faiss](https://github.com/facebookresearch/faiss/) for the dense version. You can get all of these by running:\n",
    "> pip install elasticsearch  \n",
    "> pip install faiss_gpu  \n",
    "> pip install nlp  \n",
    "> pip install transformers  \n",
    ">  \n",
    "> wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.7.1-linux-x86_64.tar.gz  \n",
    "> tar -xzvf elasticsearch-7.7.1-linux-x86_64.tar.gz  \n",
    "\n",
    "The training relies on two datasets: [ELI5](https://arxiv.org/abs/1907.09190), a processed version of the [r/explainlikeimfive](https://www.reddit.com/r/explainlikeimfive/) subreddit, and the [Wiki40b](https://www.aclweb.org/anthology/2020.lrec-1.297/) Wikipedia image. You can download both using the [🤗nlp](https://github.com/huggingface/nlp) linrary with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nlp\n",
    "eli5 = nlp.load_dataset('eli5')\n",
    "wiki40b_snippets = nlp.load_dataset('wiki_snippets', name='wiki40b_en_100_0')['train']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, all of the useful methods used in this notebook are compiled in the [lfqa_utils.py](https://github.com/huggingface/longform-qa/lfqa_utils.py) script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lfqa_utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='reddit_biases'></a>\n",
    "### 1.b - Note on Data and Biases\n",
    "\n",
    "Before we go any further, let us take a moment to talk about the provenance of our training data. While Reddit hosts a number of thriving communities with high quality discussions, it is also widely known to have corners where sexism, hate, and harassment are significant issues. See for example the [recent post from Reddit founder u/spez](https://www.reddit.com/r/announcements/comments/gxas21/upcoming_changes_to_our_content_policy_our_board/) outlining some of the ways he thinks the website's historical policies have been responsible for this problem, [Adrienne Massanari's 2015 article on GamerGate](https://www.researchgate.net/publication/283848479_Gamergate_and_The_Fappening_How_Reddit's_algorithm_governance_and_culture_support_toxic_technocultures) and follow-up works, or a [2019 Wired article on misogyny on Reddit](https://www.wired.com/story/misogyny-reddit-research/).\n",
    "\n",
    "While there has been some recent work in the NLP community on *de-biasing* models (e.g. [Black is to Criminal as Caucasian is to Police:\n",
    "Detecting and Removing Multiclass Bias in Word Embeddings](https://arxiv.org/abs/1904.04047) for word embeddings trained specifically on Reddit data), this problem is far from solved, and the likelihood that a trained model might learn the biases present in the data remains a significant concern.\n",
    "\n",
    "As mentioned above, the magnitude of the problem depends on the specific communities/subreddits. This work uses data from [r/explainlikeimfive](https://www.reddit.com/r/explainlikeimfive/), and the `nlp` library also gives access to examples from [r/askscience](https://www.reddit.com/r/askscience/), and [r/AskHistorians](https://www.reddit.com/r/AskHistorians/). There are some encouraging signs for all of these communities: [r/explainlikeimfive](https://www.reddit.com/r/explainlikeimfive/) and [r/askscience](https://www.reddit.com/r/askscience/) have similar structures and purposes, and [r/askscience](https://www.reddit.com/r/askscience/) was found in 2015 to show medium supportiveness and very low toxicity when compared to other subreddits (see a [hackerfall post](https://hackerfall.com/story/study-and-interactive-visualization-of-toxicity-in), [thecut.com write-up](https://www.thecut.com/2015/03/interactive-chart-of-reddits-toxicity.html) and supporting [data](https://chart-studio.plotly.com/~bsbell21/210/toxicity-vs-supportiveness-by-subreddit/#data)). Meanwhile, the [r/AskHistorians rules](https://www.reddit.com/r/AskHistorians/wiki/rules) mention that the admins will not tolerate \"*racism, sexism, or any other forms of bigotry*\".\n",
    "\n",
    "This is obviously not enough to exonerate the model (the pre-training step, for example, raises its own questions on that topic), and there is still a lot of interesting work to do to be able to quantify the biases in a conditional text generation model. One thing you can do to help: if you find any particularly egregious answers provided by the model when using the demo, or want to work on this research question please send a DM to [@YJernite on Twitter](https://twitter.com/YJernite)!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='task_description'></a>\n",
    "# 2. Task and Data Description\n",
    "\n",
    "Let's recap: we are interested in the task of Long Form Question Answering. As in other Question Answering tasks, the model is presented with a question, and is required to generate a natural language answer. Whereas a majority of QA datasets contain mostly **factoid** questions, where the answer, such as a date or the name of a single entity, can be expressed in a few words or single sentence, Long Form QA focuses on questions which call for an **explanation** consisting of a few sentences or a few paragraphs.\n",
    "\n",
    "In order to teach a model to answer such questions, we use questions and answers written by Reddit users. Note that the `nlp.load_dataset` command above actually downloaded questions and their associated answers from the [r/explainlikeimfive](https://www.reddit.com/r/explainlikeimfive/), [r/askscience](https://www.reddit.com/r/askscience/), and [r/AskHistorians](https://www.reddit.com/r/AskHistorians/) subreddits. We focus here on the **ELI5/explainlikeimfive** part to train the system, as these examples tend to be a little simpler.  \n",
    "\n",
    "Let's look at one item from the test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'q_id': '8houtx',\n",
       " 'title': 'Why does water heated to room temperature feel colder than the air around it?',\n",
       " 'selftext': '',\n",
       " 'document': '',\n",
       " 'subreddit': 'explainlikeimfive',\n",
       " 'answers': {'a_id': ['dylcnfk', 'dylcj49'],\n",
       "  'text': [\"Water transfers heat more efficiently than air. When something feels cold it's because heat is being transferred from your skin to whatever you're touching. Since water absorbs the heat more readily than air, it feels colder.\",\n",
       "   \"Air isn't as good at transferring heat compared to something like water or steel (sit on a room temperature steel bench vs. a room temperature wooden bench, and the steel one will feel more cold).\\n\\nWhen you feel cold, what you're feeling is heat being transferred out of you.  If there is no breeze, you feel a certain way.  If there's a breeze, you will get colder faster (because the moving air is pulling the heat away from you), and if you get into water, its quite good at pulling heat from you.   Get out of the water and have a breeze blow on you while you're wet, all of the water starts evaporating, pulling even more heat from you.\"],\n",
       "  'score': [5, 2]},\n",
       " 'title_urls': {'url': []},\n",
       " 'selftext_urls': {'url': []},\n",
       " 'answers_urls': {'url': []}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eli5['test_eli5'][12345]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So here we have the question:\n",
    "> Why does water heated to room temperature feel colder than the air around it?  \n",
    "\n",
    "\n",
    "This definitely requires a multi-step explanation: no single phrase can sum up all of the information we are looking for. Here are the answers that were given on ELI5, and were given scores of +5 and +2 respectively by Reddit users:\n",
    "> 1. Water transfers heat more efficiently than air. When something feels cold it's because heat is being transferred from your skin to whatever you're touching. Since water absorbs the heat more readily than air, it feels colder.  \n",
    "\n",
    "> 2. Air isn't as good at transferring heat compared to something like water or steel (sit on a room temperature steel bench vs. a room temperature wooden bench, and the steel one will feel more cold). When you feel cold, what you're feeling is heat being transferred out of you. If there is no breeze, you feel a certain way.  If there's a breeze, you will get colder faster (because the moving air is pulling the heat away from you), and if you get into water, its quite good at pulling heat from you. Get out of the water and have a breeze blow on you while you're wet, all of the water starts evaporating, pulling even more heat from you.  \n",
    "\n",
    "First, note that in this case **we have two answers** which broadly describe the same phenomenon: the first one is scored higher because it is more succint and to the point. This example already illustrates one important feature of the LFQA task: **there are usually several valid ways to answer a given question.** Of the 272K examples in the ELI5 training set, nearly two thirds (167K) have at least two answers. We'll need to keep this in mind when training and evaluation of the model.  \n",
    "\n",
    "Secondly, we need to give our model access to the information that is expressed in both these answers. Recently released models have been shown to include a significant amount of world knowledge in their parameters without the need of any external knowledge at all (see e.g. the [Closed-book QA performance of the T5 model](https://arxiv.org/abs/2002.08910)). There are several advantages to giving the model explicit access to information in text form however. First, a larger number of parameters in a model implies a larger computational cost. Secondly, getting information from a text database allows us to easily update the model's knowledge without having to re-train its parameters.\n",
    "\n",
    "--- \n",
    "\n",
    "<img src=\"images/ELI5animation.gif\" width=\"750\" align=\"center\"/>  \n",
    "\n",
    "---\n",
    "\n",
    "<center> Overview of the full question answering process.</center>\n",
    "<center> First, the Document Retriever selects a set of passages from Wikipedia that have information relevant to the question.</center>\n",
    "<center> Then, the Answer Generation Model reads the concatenation of the question and retrieverd passages, and writes out the answer.</center>\n",
    "\n",
    "---  \n",
    "Here, we choose to give the model access to Wikipedia text. Full Wikipedia articles are typically too long for most current models to handle, and notable exceptions like the [Reformer](https://arxiv.org/abs/2001.04451) or [Longformer](https://arxiv.org/abs/2004.05150) architectures unfortunately do not yet have pre-trained sequence-to-sequence variants. Thus, we follow previous work in splitting Wikipedia articles into disjoint snippets of 100 words, and keep track of the title of the article and sections a snippet came from. Here's how you can get a pre-processed Wiki40b version split into 100-word passages with the `nlp` library, and an example snippet which has some of the information we're looking for (\"*little conduction would occur since air is a poor conductor of heat*\"):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_id': '{\"nlp_id\": 1665419, \"wiki_id\": \"Q179635\", \"sp\": 12, \"sc\": 653, \"ep\": 12, \"ec\": 1223}',\n",
       " 'nlp_id': 1665419,\n",
       " 'wiki_id': 'Q179635',\n",
       " 'start_paragraph': 12,\n",
       " 'start_character': 653,\n",
       " 'end_paragraph': 12,\n",
       " 'end_character': 1223,\n",
       " 'article_title': 'Heat transfer',\n",
       " 'section_title': 'Conduction',\n",
       " 'passage_text': 'from one place to another place without the movement of particles is called conduction, such as when placing a hand on a cold glass of water - heat is conducted from the warm skin to the cold glass, but if the hand is held a few inches from the glass, little conduction would occur since air is a poor conductor of heat. Steady state conduction is an idealized model of conduction that happens when the temperature difference driving the conduction is constant, so that after a time, the spatial distribution of temperatures in the conducting object does not change any'}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wiki40b_snippets[8991855]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next two sections, we show how we can use either a [sparse retriever](#elasticsearch) or a [trained dense retriever](#dense_retrieval) to automatically find relevant snippets for a question."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='elasticsearch'></a>\n",
    "# 3. Sparse Retrieval: Making Support Documents with ElasticSearch\n",
    "\n",
    "In this section, we show how to use either such a \"classical\" Information Retrieval (IR) system based on **sparse** word matching with [ElasticSearch](https://www.elastic.co/elasticsearch/), an extremely popular and efficient search engine that can be used for finding documents that match a given query based on word overlap.  \n",
    "\n",
    "Specifically, ElasticSearch provides a convenient way to index documents so they can easily be queried for nearest neighbor search using the [BM25](https://en.wikipedia.org/wiki/Okapi_BM25) similarity function (which relies on TF-IDF weighting of words). While this word-matching based approach has obvious limitations, such as failing to take synonyms and sometimes grammatical variation into account, it does pretty well overall and has only recently been overtaken by embedding-based systems for Wikipedia-based Open-Domain QA tasks.\n",
    "\n",
    "In order to use ElasticSearch, you will first need to launch a server. In a different window, run:\n",
    "> ./elasticsearch-7.7.0/bin/elasticsearch\n",
    "\n",
    "By default, your ElasticSearch server will be listening on `localhost` port `9200`. To connect to it run: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "es_client = Elasticsearch([{'host': 'localhost', 'port': '9200'}])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `eli5_utils.py` script provides utilities to create (`make_es_index_snippets`) and query (`query_es_index`) an ElasticSearch index from within Python.\n",
    "\n",
    "The main implementation details are:  \n",
    "1. We index the article title, section title, and text of each of the passages for BM25 passages. These choices are implemented in the `index_config` variable:  \n",
    "```python\n",
    "    index_config = {\n",
    "      \"settings\": {\n",
    "        \"number_of_shards\": 1,\n",
    "      },\n",
    "      \"mappings\": {\n",
    "        \"properties\": {\n",
    "          \"article_title\": {\"type\": \"text\", \"analyzer\": \"standard\", \"similarity\": \"BM25\"},\n",
    "          \"section_title\": {\"type\": \"text\", \"analyzer\": \"standard\", \"similarity\": \"BM25\"},\n",
    "          \"passage_text\": {\"type\": \"text\", \"analyzer\": \"standard\", \"similarity\": \"BM25\"}\n",
    "        }\n",
    "      }\n",
    "    }\n",
    "```\n",
    "2. To query the index, we found it useful to add a few task-dependent stop-words. The query text is then compared to all of the indexed fields, giving more weight to the passage text:\n",
    "```python\n",
    "    banned = ['how', 'why', 'what', 'where', 'which', 'do', 'does', 'is', '?', 'eli5', 'eli5:']\n",
    "    q = ' '.join([w for w in q.split() if w not in banned])\n",
    "    response = es_client.search(\n",
    "        index = index_name,\n",
    "        body = {\n",
    "            \"query\": {\n",
    "                \"multi_match\": {\n",
    "                    \"query\": q,\n",
    "                    \"fields\": [\"article_title\", \"section_title\", \"passage_text^2\"],\n",
    "                    \"type\": \"cross_fields\",\n",
    "                }\n",
    "            },\n",
    "            \"size\": n_results,\n",
    "        }\n",
    "    )\n",
    "```  \n",
    "\n",
    "Here's the command to create the index, it should take one to three hours depending on your system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not es_client.indices.exists('wiki40b_snippets_100w'):\n",
    "    make_es_index_snippets(es_client, wiki40b_snippets, index_name='wiki40b_snippets_100w')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's test the ElasticSearch retriever with our running example ELI5 question about skin-to-water heat transfer by returning the 10 best candidate passages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row0_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row0_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row0_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row1_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row1_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row1_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row2_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row2_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row2_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row3_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row3_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row3_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row4_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row4_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row4_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row5_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row5_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row5_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row6_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row6_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row6_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row7_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row7_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row7_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row8_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row8_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row8_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row9_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row9_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row9_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row10_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row10_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_e6a67934_abe6_11ea_bac0_cb508f4c6806row10_col2 {\n",
       "            text-align:  left;\n",
       "        }</style><table id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >Article</th>        <th class=\"col_heading level0 col1\" >Sections</th>        <th class=\"col_heading level0 col2\" >Text</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row0_col0\" class=\"data row0 col0\" >---</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row0_col1\" class=\"data row0 col1\" >---</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row0_col2\" class=\"data row0 col2\" >--- Why does water heated to room temperature feel colder than the air around it?</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row1\" class=\"row_heading level0 row1\" >1</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row1_col0\" class=\"data row1 col0\" >Salt fingering</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row1_col1\" class=\"data row1 col1\" >Salt fingering</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row1_col2\" class=\"data row1 col2\" >Salt fingering Salt fingering is a mixing process that occurs when relatively warm, salty water overlies relatively colder, fresher water. It is driven by the fact that heated water diffuses more readily than salty water. A small parcel of warm, salty water sinking downwards into a colder, fresher region will lose its heat before losing its salt, making the parcel of water increasingly denser than the water around it and sinking further. Likewise, a small parcel of colder, fresher water will be displaced upwards and gain heat by diffusion from surrounding water, which will then make it lighter than the</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row2\" class=\"row_heading level0 row2\" >2</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row2_col0\" class=\"data row2 col0\" >Solar water heating</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row2_col1\" class=\"data row2 col1\" >Flat plate & Evacuated tube</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row2_col2\" class=\"data row2 col2\" >protected by a glass panel. Consequently, these types of collectors are much less efficient when water temperature exceeds ambient air temperatures. For pool heating applications, the water to be heated is often colder than the ambient roof temperature, at which point the lack of thermal insulation allows additional heat to be drawn from the surrounding environment. Evacuated tube Evacuated tube collectors (ETC) are a way to reduce the heat loss, inherent in flat plates. Since heat loss due to convection cannot cross a vacuum, it forms an efficient isolation mechanism to keep heat inside the collector pipes. Since two flat</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row3\" class=\"row_heading level0 row3\" >3</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row3_col0\" class=\"data row3 col0\" >Humidifier</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row3_col1\" class=\"data row3 col1\" >Fixed-installation humidifiers & Problems</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row3_col2\" class=\"data row3 col2\" >to damage from overly dry air. In colder months, they may provide modest energy savings, since as humidity increases, occupants may feel warm at a lower temperature.\n",
       "Bypass humidifiers are connected between the heated and cold air return ducts, using the pressure difference between these ducts to cause some heated air to make a bypass through the humidifier and return to the furnace.\n",
       "Any humidifiers should usually be disabled during the summer months if air conditioning is used; air conditioners partially function to reducing indoor humidity, and having a humidifier continue to operate will waste significant amounts of energy. Problems The USEPA</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row4\" class=\"row_heading level0 row4\" >4</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row4_col0\" class=\"data row4 col0\" >Drake Landing Solar Community</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row4_col1\" class=\"data row4 col1\" >How it works & Energy centre</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row4_col2\" class=\"data row4 col2\" >the short-term storage tanks in the Energy Centre to be heated again in order to complete the circuit. During colder months the water from the BTES passes back to the short-term storage tank and is then directed to each home. Similar to a hot water tank, the heated water goes through a heat exchanger that blows air across the warm fan coil. Heat travels from the water to the air and is directed through the house via ductwork. When the temperature reaches that said on the thermostat, an automatic valve shuts off the heat transfer unit. Energy centre The Energy</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row5\" class=\"row_heading level0 row5\" >5</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row5_col0\" class=\"data row5 col0\" >Diamond dust</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row5_col1\" class=\"data row5 col1\" >Characteristics & Formation</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row5_col2\" class=\"data row5 col2\" >it looks like many tiny diamonds are flashing in the air. Formation These ice crystals usually form when a temperature inversion is present at the surface and the warmer air above the ground mixes with the colder air near the surface. Since warmer air frequently contains more water vapor than colder air, this mixing will usually also transport water vapor into the air near the surface, causing the relative humidity of the near-surface air to increase. If the relative humidity increase near the surface is large enough then ice crystals may form.\n",
       "To form diamond dust the temperature must be below</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row6\" class=\"row_heading level0 row6\" >6</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row6_col0\" class=\"data row6 col0\" >Effects of global warming on oceans</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row6_col1\" class=\"data row6 col1\" >Ocean currents</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row6_col2\" class=\"data row6 col2\" >changing latitudes of our planet. As the atmosphere is warmed nearest the equator, the hot air at the surface of our planet is heated, causing it to rise and draw in cooler air to take its place, creating what is known as circulation cells. This ultimately causes the air to be significantly colder near the poles than at the equator.\n",
       "Wind patterns associated with these circulation cells drive surface currents which push the surface water to the higher latitudes where the air is colder. This cools the water down enough to where it is capable of dissolving more gasses and minerals,</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row7\" class=\"row_heading level0 row7\" >7</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row7_col0\" class=\"data row7 col0\" >Mesoscale convective system</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row7_col1\" class=\"data row7 col1\" >Lake-effect snow</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row7_col2\" class=\"data row7 col2\" >for lake-effect rain or snow to form, the air moving across the lake must be significantly cooler than the surface air (which is likely to be near the temperature of the water surface). Specifically, the air temperature at the altitude where the air pressure is 850 millibars (or 1.5 kilometres (0.93 mi) altitude) should be 13 °C (24 °F) lower than the temperature of the air at the surface.  Lake-effect occurring when the air at 850 millibars is 25 °C (45 °F) colder than the water temperature can produce thundersnow, snow showers accompanied by lightning and thunder (due to the larger amount of energy</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row8\" class=\"row_heading level0 row8\" >8</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row8_col0\" class=\"data row8 col0\" >Thermal comfort</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row8_col1\" class=\"data row8 col1\" >Interplay of temperature and humidity</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row8_col2\" class=\"data row8 col2\" >such as the heat index.\n",
       "For lower temperatures, a related interplay was identified only qualitatively:\n",
       "High humidity and low temperatures cause the air to feel chilly.\n",
       "Cold air with high relative humidity \"feels\" colder than dry air of the same temperature because high humidity in cold weather increases the conduction of heat from the body.\n",
       "There has been controversy over why damp cold air feels colder than dry cold air. Some believe it is because when the humidity is high, our skin and clothing become moist and are better conductors of heat, so there is more cooling by conduction.\n",
       "For more recent data look for</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row9\" class=\"row_heading level0 row9\" >9</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row9_col0\" class=\"data row9 col0\" >Honyaki</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row9_col1\" class=\"data row9 col1\" >Traditional process</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row9_col2\" class=\"data row9 col2\" >heat. The quench water or oil is prepared and brought to the right temperature. The forge is heated. Lights are turned off and the room shut from the outside. Once ready, the blade is buried and shuffled around in the charcoal and when it reaches the correct temperature it is thrust into water and moved forward and back (so as to prevent lateral distortion) and then after a couple seconds side to side. The knife could also be brought up slightly above temperature and then held in the room to the correct temperature before quench. The quench could be interrupted</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806level0_row10\" class=\"row_heading level0 row10\" >10</th>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row10_col0\" class=\"data row10 col0\" >Greywell Tunnel</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row10_col1\" class=\"data row10 col1\" >SSSI</td>\n",
       "                        <td id=\"T_e6a67934_abe6_11ea_bac0_cb508f4c6806row10_col2\" class=\"data row10 col2\" >creates an ideal micro-climate for the bats, which is maintained at around 10 °C (50 °F) all year. When the temperature outside the tunnel is colder than this, cold air flows into the bottom of the tunnel where it is warmed by the water, and warmer air flows out along the top of the tunnel. During the summer the air flow is reversed, with warm air flowing into the top of the tunnel and being cooled as it flows back out over the water. By 2006, there were some 12,500 bats roosting in the tunnel, which included the largest known colony of</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f0c92a4f610>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "question = eli5['test_eli5'][12345]['title']\n",
    "doc, res_list = query_es_index(question, es_client, index_name='wiki40b_snippets_100w')\n",
    "\n",
    "df = pd.DataFrame({\n",
    "    'Article': ['---'] + [res['article_title'] for res in res_list],\n",
    "    'Sections': ['---'] + [res['section_title'] if res['section_title'].strip() != '' else res['article_title']\n",
    "                 for res in res_list],\n",
    "    'Text': ['--- ' + question] + [res['passage_text'] for res in res_list],\n",
    "})\n",
    "df.style.set_properties(**{'text-align': 'left'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can immediately see both the strengths and limitations of this approach. The system manages to retrieve documents that are all broadly on topic, mentioning some combination of *water*, *air*, *relative temperature*, and *temperature transfer*. In spite of this, only example 8 ends up containing information that is actually relevant to the question:\n",
    "> Cold air with high relative humidity \"feels\" colder than dry air of the same temperature because high humidity in cold weather increases the conduction of heat from the body.  \n",
    "\n",
    "We got lucky this time, but this passage could as easily have been ranked 11th and not been included in the support document we provide to the answer generation system. As it is, the model will have to sort through mostly off-topic information to find this sentence when reading the resulting supporting document."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='dense_retrieval'></a>\n",
    "# 4. Retrieving Support Documents with an ELI5-Trained Dense Model\n",
    "\n",
    "The sparse retriever works by finding passages which feature the words from the query. However, it has no way to know *a priori* which of these words are more important in context, and seems to struggle with understanding the central theme of the query (human-perceived temperature).\n",
    "\n",
    "Thankfully, some recent works have taken advantage of advances in pre-trained contextual word representations to solve this problem. Models such as [DPR](https://arxiv.org/abs/2004.04906) or [REALM](https://arxiv.org/abs/2002.08909) for example learn to compute a vector representation of the query, as well as vector representations of Wikipedia passages in such a way that the passages that best answers a question maximize the dot product between the two representations. Retrieval is then reduced to a Maximum Inner Product Search, which can be executed efficiently using systems like [FAISS](https://github.com/facebookresearch/faiss).\n",
    "\n",
    "These successes are very encouraging for our Open-Domain Long Form QA application. However, our task and setup do not quite meet the requirements of either of either of these approaches. On the one hand, the [DPR](https://arxiv.org/abs/2004.04906) system is trained using gold passage annotations: most major QA dataset tell the system which Wikipedia passage contains the answer. Unfortunately, we do not have such annotations for the ELI5 data. On the other hand, while [REALM](https://arxiv.org/abs/2002.08909) is trained without passage supervision, it requires a pretty expensive pre-training step with an [Inverse Cloze Task](https://arxiv.org/abs/1906.00300) (100,000 steps with batch size 4096), and the ability to re-compute the embeddings of all Wikipedia passages regularly during training.\n",
    "\n",
    "In order to train a similar dense retrieval system at reduced cost without having access to gold passage annotation, we will have to **take advantage of another unique feature of our dataset**, namely the fact that the long form answers are quite similar in style to the Wikipedia passages we want to index. Our hypothesis then is that if we train a system to embed the questions and answers in our dataset in a way that allows us to easily match questions to answers, then using the answer embedder on Wikipedia passages should allow us to similarly match questions to supporting evidence from Wikipedia.\n",
    "\n",
    "<a id='dense_train'></a>\n",
    "### 4.a - Contrastive Training with ELI5 In-Batch Negatives\n",
    "\n",
    "As mentioned above, we want to train a system to produce question and answer embeddings, such that the dot product between the representation of a question and any of its answers is greater than between it and answers of all of the other questions in the dataset.  \n",
    "\n",
    "Unfortunately, actually comparing all questions to all answers before taking every single gradient step is computationally prohibitive: instead, we follow previous work in simply processing medium to large batches of question-answer pairs, and making sure that the dot product of a question with its answer is larger than with all other answers in the batch, and *vice versa*.  \n",
    "\n",
    "We use a cross-entropy loss for the multinomial distribution over all of the answers (or questions) in a batch, and make use of [PyTorch gradient checkpointing](https://pytorch.org/docs/stable/checkpoint.html) to be able to use large batches with limited GPU memory: you can find all implementation details in the `RetrievalQAEmbedder` class in `eli5_utils.py`.\n",
    "\n",
    "---  \n",
    "\n",
    "<img src=\"images/ELI5contrastive.svg\" width=\"700\" align=\"center\"/>  \n",
    "\n",
    "---  \n",
    "\n",
    "<center> To train the retriever, we show the model batches of 512 question-answer pairs.</center>\n",
    "<center> The model needs to ensure that the embedding of each question in the batch is closer to the embedding</center>\n",
    "<center> of its corresponding answer than to the embedding of any other answer in the batch.</center>\n",
    "\n",
    "---  \n",
    "\n",
    "\n",
    "We use a single BERT-style pre-trained model to embed the questions and answers, and learn different projection matrices to bring both representations down to dimension 128: the projection matrices are trained from scratch as the sentence embedding model is fine-tuned. We found that the 8-layer distilled version of BERT from the [Well-Read Students Learn Better paper](https://arxiv.org/abs/1908.08962) performed as well or better as full BERT for a notable gain in computation speed: if you want an even faster model, that work provides pre-trained models spanning the full range of computation/accuracy trade-offs.\n",
    "\n",
    "The model can than be trained with the following code: with batch size 32/512 on a single 16GB GPU, you can run 10 training epochs in under 6 hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training arguments\n",
    "class ArgumentsQAR():\n",
    "    def __init__(self):\n",
    "        self.batch_size = 512\n",
    "        self.max_length = 128\n",
    "        self.checkpoint_batch_size = 32\n",
    "        self.print_freq = 100\n",
    "        self.pretrained_model_name = \"google/bert_uncased_L-8_H-768_A-12\"\n",
    "        self.model_save_name = \"retriever_models/eli5_retriever_model_l-8_h-768_b-512-512\"\n",
    "        self.learning_rate = 2e-4\n",
    "        self.num_epochs = 10\n",
    "\n",
    "qar_args = ArgumentsQAR()\n",
    "\n",
    "# prepare torch Dataset objects\n",
    "qar_train_dset = ELI5DatasetQARetriver(eli5['train_eli5'], training=True)\n",
    "qar_valid_dset = ELI5DatasetQARetriver(eli5['validation_eli5'], training=False)\n",
    "\n",
    "# load pre-trained BERT and make model\n",
    "qar_tokenizer, qar_model = make_qa_retriever_model(\n",
    "        model_name=qar_args.pretrained_model_name,\n",
    "        from_file=None,\n",
    "        device=\"cuda:0\"\n",
    ")\n",
    "\n",
    "# train the model\n",
    "train_qa_retriever(qar_model, qar_tokenizer, qar_train_dset, qar_valid_dset, qar_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you don't want to train the model yourself, you can also download trained weights from the [Hugging Face model repository](https://huggingface.co/models) with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "qar_tokenizer = AutoTokenizer.from_pretrained('yjernite/retribert-base-uncased')\n",
    "qar_model = AutoModel.from_pretrained('yjernite/retribert-base-uncased').to('cuda:1')\n",
    "_ = qar_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the model is trained, it can be used to compute passage embeddings for all Wikipedia snippets. The `make_qa_dense_index` method takes advantage of `numpy` memory-mapping, so embeddings are written directly to disk. Again with a single GPU, computing the full set of passage embeddings should take about 18 hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.isfile('wiki40b_passages_reps_32_l-8_h-768_b-512-512.dat'):\n",
    "    make_qa_dense_index(\n",
    "        qar_model, qar_tokenizer, wiki40b_snippets, device='cuda:0',\n",
    "        index_name='wiki40b_passages_reps_32_l-8_h-768_b-512-512.dat'\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='dense_use'></a>\n",
    "### 4.b -  Using the Trained Dense Retriever and Wikipedia Index\n",
    "\n",
    "Now that we have trained our model to compute query and answer embeddings and used it to compute passage embeddings for all our Wikipedia snippets, let's see whether it can actually find supporting evidence for a new question. Recall the the two steps to using the dense retriever: we first compute an embedding for a new question, then do Max Inner Product Search with the pre-computed passage representations.\n",
    "\n",
    "---  \n",
    "\n",
    "<img src=\"images/ELI5wiki_index.svg\" width=\"600\" align=\"center\"/>  \n",
    "\n",
    "---  \n",
    "\n",
    "<center> At test time, the Retriever Model encodes the question, and compares its embedding to the pre-computed representation of</center>\n",
    "<center>  all the Wikipedia passages. The ten passages with the closest embedding are returned to create the support document.</center>\n",
    "\n",
    "---  \n",
    "\n",
    "The MIPS part can be executed efficiently with the `faiss` library. Additionally, since we computed 128-dimensional passage embeddings, the whole of the representations fits on a GPU, making retrieval even faster. We can create the `faiss_gpu` index with the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "faiss_res = faiss.StandardGpuResources()\n",
    "wiki40b_passage_reps = np.memmap(\n",
    "            'wiki40b_passages_reps_32_l-8_h-768_b-512-512.dat',\n",
    "            dtype='float32', mode='r',\n",
    "            shape=(wiki40b_snippets.num_rows, 128)\n",
    ")\n",
    "\n",
    "wiki40b_index_flat = faiss.IndexFlatIP(128)\n",
    "wiki40b_gpu_index = faiss.index_cpu_to_gpu(faiss_res, 1, wiki40b_index_flat)\n",
    "wiki40b_gpu_index.add(wiki40b_passage_reps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use the `query_qa_dense_index` function to query the dense index for our running example question about perceived temperature:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row0_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row0_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row0_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row1_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row1_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row1_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row2_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row2_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row2_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row3_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row3_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row3_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row4_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row4_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row4_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row5_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row5_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row5_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row6_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row6_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row6_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row7_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row7_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row7_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row8_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row8_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row8_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row9_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row9_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row9_col2 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row10_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row10_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_86da2ed6_acec_11ea_ba32_1b67035d3076row10_col2 {\n",
       "            text-align:  left;\n",
       "        }</style><table id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >Article</th>        <th class=\"col_heading level0 col1\" >Sections</th>        <th class=\"col_heading level0 col2\" >Text</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row0_col0\" class=\"data row0 col0\" >---</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row0_col1\" class=\"data row0 col1\" >---</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row0_col2\" class=\"data row0 col2\" >--- Why does water heated to room temperature feel colder than the air around it?</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row1\" class=\"row_heading level0 row1\" >1</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row1_col0\" class=\"data row1 col0\" >Heat transfer</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row1_col1\" class=\"data row1 col1\" >Heat transfer in the human body & Evaporative cooling</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row1_col2\" class=\"data row1 col2\" >when the skin is completely wet. The body continuously loses water by evaporation but the most significant amount of heat loss occurs during periods of increased physical activity. Evaporative cooling Evaporative cooling happens when water vapor is added to the surrounding air. The energy needed to evaporate the water is taken from the air in the form of sensible heat and converted into latent heat, while the air remains at a constant enthalpy. Latent heat describes the amount of heat that is needed to evaporate the liquid; this heat comes from the liquid itself and the surrounding gas and surfaces.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row2\" class=\"row_heading level0 row2\" >2</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row2_col0\" class=\"data row2 col0\" >Johan Sandström</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row2_col1\" class=\"data row2 col1\" >Sandström  Theorem</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row2_col2\" class=\"data row2 col2\" >at greater pressures. There is an ambiguity, however, as to the meaning of the terms 'heating' and 'cooling' in Sandstrom's theorem. So far, heating and cooling has always been interpreted in the literature as being associated with 'surface heating' and 'surface cooling' respectively.\n",
       "In real fluids, however, molecular and turbulent diffusion always cause internal heating/cooling even in absence of external heating/cooling, as long as the temperature of the fluid considered is non-uniform. As is well-known, molecular and turbulent diffusion tends to relax the system toward thermodynamic equilibrium, i.e., toward an isothermal state, which for a statically stable fluid, will warm up</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row3\" class=\"row_heading level0 row3\" >3</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row3_col0\" class=\"data row3 col0\" >Thermal equilibrium</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row3_col1\" class=\"data row3 col1\" >Bodies prepared with separately uniform temperatures, then put into purely thermal communication with each other</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row3_col2\" class=\"data row3 col2\" >are not in a relation of thermal equilibrium, heat will flow from the hotter to the colder, by whatever pathway, conductive or radiative, is available, and this flow will continue until thermal equilibrium is reached and then they will have the same temperature.\n",
       "One form of thermal equilibrium is radiative exchange equilibrium. Two bodies, each with its own uniform temperature, in solely radiative connection, no matter how far apart, or what partially obstructive, reflective, or refractive, obstacles lie in their path of radiative exchange, not moving relative to one another, will exchange thermal radiation, in net the hotter transferring energy to</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row4\" class=\"row_heading level0 row4\" >4</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row4_col0\" class=\"data row4 col0\" >Evaporative cooler</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row4_col1\" class=\"data row4 col1\" >Physical principles</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row4_col2\" class=\"data row4 col2\" >air condition and moving along a line of constant enthalpy toward a state of higher humidity.\n",
       "A simple example of natural evaporative cooling is perspiration, or sweat, secreted by the body, evaporation of which cools the body. The amount of heat transfer depends on the evaporation rate, however for each kilogram of water vaporized 2,257 kJ of energy (about 890 BTU per pound of pure water, at 95 °F (35 °C)) are transferred. The evaporation rate depends on the temperature and humidity of the air, which is why sweat accumulates more on humid days, as it does not evaporate fast enough.\n",
       "Vapor-compression refrigeration uses evaporative cooling,</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row5\" class=\"row_heading level0 row5\" >5</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row5_col0\" class=\"data row5 col0\" >Thermal contact conductance</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row5_col1\" class=\"data row5 col1\" >Factors influencing contact conductance & Contact pressure</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row5_col2\" class=\"data row5 col2\" >Thermal contact conductance In physics, thermal contact conductance is the study of heat conduction between solid bodies in thermal contact. The thermal contact conductance coefficient, , is a property indicating the thermal conductivity, or ability to conduct heat, between two bodies in contact. The inverse of this property is termed thermal contact resistance. Factors influencing contact conductance Thermal contact conductance is a complicated phenomenon, influenced by many factors. Experience shows that the most important ones are as follows: Contact pressure For thermal transport between two contacting bodies, such as particles in a granular medium, the contact pressure is the factor</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row6\" class=\"row_heading level0 row6\" >6</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row6_col0\" class=\"data row6 col0\" >Thermodynamic temperature</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row6_col1\" class=\"data row6 col1\" >The heat of phase changes</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row6_col2\" class=\"data row6 col2\" >to completely boil or vaporize water (what is known as enthalpy of vaporization) is roughly 540 times that required for a one-degree increase.\n",
       "Water's sizable enthalpy of vaporization is why one's skin can be burned so quickly as steam condenses on it (heading from red to green in Fig. 7 above). In the opposite direction, this is why one's skin feels cool as liquid water on it evaporates (a process that occurs at a sub-ambient wet-bulb temperature that is dependent on relative humidity). Water's highly energetic enthalpy of vaporization is also an important factor underlying why solar pool covers (floating, insulated blankets that</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row7\" class=\"row_heading level0 row7\" >7</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row7_col0\" class=\"data row7 col0\" >Temperature</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row7_col1\" class=\"data row7 col1\" >Local thermodynamic equilibrium & Bodies in thermodynamic equilibrium</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row7_col2\" class=\"data row7 col2\" >and this is because temperature is an intensive variable. Bodies in thermodynamic equilibrium For experimental physics, hotness means that, when comparing any two given bodies in their respective separate thermodynamic equilibria, any two suitably given empirical thermometers with numerical scale readings will agree as to which is the hotter of the two given bodies, or that they have the same temperature. This does not require the two thermometers to have a linear relation between their numerical scale readings, but it does require that the relation between their numerical readings shall be strictly monotonic. A definite sense of greater hotness can</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row8\" class=\"row_heading level0 row8\" >8</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row8_col0\" class=\"data row8 col0\" >Latent heat</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row8_col1\" class=\"data row8 col1\" >Usage</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row8_col2\" class=\"data row8 col2\" >phase of atmospheric or ocean water,  vaporization, condensation, freezing or melting, whereas sensible heat is energy transferred that is evident in change of the temperature of the atmosphere or ocean, or ice, without those phase changes, though it is associated with changes of pressure and volume.\n",
       "The original usage of the term, as introduced by Black, was applied to systems that were intentionally held at constant temperature. Such usage referred to latent heat of expansion and several other related latent heats. These latent heats are defined independently of the conceptual framework of thermodynamics.\n",
       "When a body is heated at constant temperature</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row9\" class=\"row_heading level0 row9\" >9</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row9_col0\" class=\"data row9 col0\" >Sensible heat</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row9_col1\" class=\"data row9 col1\" >Sensible heat</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row9_col2\" class=\"data row9 col2\" >Sensible heat Sensible heat is heat exchanged by a body or thermodynamic system in which the exchange of heat changes the temperature of the body or system, and some macroscopic variables of the body or system, but leaves unchanged certain other macroscopic variables of the body or system, such as volume or pressure.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076level0_row10\" class=\"row_heading level0 row10\" >10</th>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row10_col0\" class=\"data row10 col0\" >Heat transfer</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row10_col1\" class=\"data row10 col1\" >Overview & Conduction</td>\n",
       "                        <td id=\"T_86da2ed6_acec_11ea_ba32_1b67035d3076row10_col2\" class=\"data row10 col2\" >changes. Conduction On a microscopic scale, heat conduction occurs as hot, rapidly moving or vibrating atoms and molecules interact with neighboring atoms and molecules, transferring some of their energy (heat) to these neighboring particles. In other words, heat is transferred by conduction when adjacent atoms vibrate against one another, or as electrons move from one atom to another. Conduction is the most significant means of heat transfer within a solid or between solid objects in thermal contact. Fluids—especially gases—are less conductive. Thermal contact conductance is the study of heat conduction between solid bodies in contact. The process of heat transfer</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7fd995420110>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "question = eli5['test_eli5'][12345]['title']\n",
    "doc, res_list = query_qa_dense_index(question, qar_model, qar_tokenizer, wiki40b_snippets, wiki40b_gpu_index, device='cuda:1')\n",
    "\n",
    "df = pd.DataFrame({\n",
    "    'Article': ['---'] + [res['article_title'] for res in res_list],\n",
    "    'Sections': ['---'] + [res['section_title'] if res['section_title'].strip() != '' else res['article_title']\n",
    "                 for res in res_list],\n",
    "    'Text': ['--- ' + question] + [res['passage_text'] for res in res_list],\n",
    "})\n",
    "df.style.set_properties(**{'text-align': 'left'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The retrieved documents are quite different from the ones returned by the sparse retrieval, with a greater focus on how water helps draw heat from a body, either through evaporation or through better conduction, which is information the model needs to answer this question.\n",
    "\n",
    "The retriever still misses out on one aspect of the query: the way the question is formulated implies that in the considered scenario the person is immersed in water rather than just wet, which makes the \"latent heat\" and evaporation arguments a little less relevant, but that's a really subtle distinction!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='dense_eval'></a>\n",
    "### 4.c -  Retriever Model Evaluation\n",
    "\n",
    "We have trained a retrieval model that *seems* to be working a little better than the traditional word-matching based approach, at least on our running example. Before we use it to actually answer questions, however, we would like to be able to get some **quantitative evaluation** of the performances of both approaches.\n",
    "\n",
    "For the retriever, we want to favor **recall** over precision: our first priority is to make sure that all of the information needed to write the answers is present in the support document. If there is unrelated information, the generation model can learn to sort it out. We measure this by computing the proportion of words in the high-scoring answers which are present in the retrieved support document. To focus on important words, we also weigh answer words by their *Inverse Document Frequency*. This gives us the following **IDF-recall** scoring function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We first select high-scoring answers (answers beyond the first must have a score of at least 3)\n",
    "test_qa_list = [(exple['title'],\n",
    "                ' '.join([a \n",
    "                          for i, (a, sc) in enumerate(zip(exple['answers']['text'], exple['answers']['score'])) \\\n",
    "                          if i == 0 or sc >= 3\n",
    "                         ]))\n",
    "                for exple in eli5['test_eli5']]\n",
    "\n",
    "# We then compute word frequencies in answer text\n",
    "answer_doc_freq = {}\n",
    "for q, a in test_qa_list:\n",
    "    for w in a.lower().split():\n",
    "        answer_doc_freq[w] = answer_doc_freq.get(w, 0) + 1\n",
    "\n",
    "# The IDF-recall function is then:\n",
    "def da_idf_recall(doc, answer):\n",
    "    d_words = dict([(w, True) for w in doc.lower().split()])\n",
    "    a_words = answer.lower().split()   \n",
    "    recall = sum([1. / math.log(1 + answer_doc_freq.get(w, 1)) for w in a_words if w in d_words]) / \\\n",
    "                sum([1. / math.log(1 + answer_doc_freq.get(w, 1)) for w in a_words])\n",
    "    return recall"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `evaluate_retriever` function in `eli5_utils.py` takes a retrieval and scoring function and computes both the average retrieval time and score of the document relative the the provided answer. Let's write some short-hand functions for the dense and sparse retrievers with our currently loaded indexes, and evaluate them on the ELI5 test set (be advised that evaluating the retriever on the full test set takes up to two hours):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "</style><table id=\"T_9934f8ca_abfe_11ea_bac0_cb508f4c6806\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >IDF-Recall</th>        <th class=\"col_heading level0 col1\" >Time/Query</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_9934f8ca_abfe_11ea_bac0_cb508f4c6806level0_row0\" class=\"row_heading level0 row0\" >Sparse</th>\n",
       "                        <td id=\"T_9934f8ca_abfe_11ea_bac0_cb508f4c6806row0_col0\" class=\"data row0 col0\" >0.3212</td>\n",
       "                        <td id=\"T_9934f8ca_abfe_11ea_bac0_cb508f4c6806row0_col1\" class=\"data row0 col1\" >0.3162</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_9934f8ca_abfe_11ea_bac0_cb508f4c6806level0_row1\" class=\"row_heading level0 row1\" >Dense</th>\n",
       "                        <td id=\"T_9934f8ca_abfe_11ea_bac0_cb508f4c6806row1_col0\" class=\"data row1 col0\" >0.3247</td>\n",
       "                        <td id=\"T_9934f8ca_abfe_11ea_bac0_cb508f4c6806row1_col1\" class=\"data row1 col1\" >0.0948</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f0993c28f90>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def dense_ret_for_eval(question, n_ret):\n",
    "    _, dense_res_list = query_qa_dense_index(\n",
    "        question, qar_model, qar_tokenizer, wiki40b_snippets, wiki40b_gpu_index, n_results=n_ret, device='cuda:1'\n",
    "    )\n",
    "    dense_doc = ' '.join([res['passage_text'] for res in dense_res_list])\n",
    "    return dense_doc\n",
    "\n",
    "def sparse_ret_for_eval(question, n_ret):\n",
    "    _, sparse_res_list = query_es_index(\n",
    "        question, es_client, index_name='wiki40b_snippets_100w', n_results=n_ret\n",
    "    )\n",
    "    sparse_doc = ' '.join([res['passage_text'] for res in sparse_res_list])\n",
    "    return sparse_doc\n",
    "\n",
    "dense_score = evaluate_retriever(test_qa_list, dense_ret_for_eval, da_idf_recall)\n",
    "sparse_score = evaluate_retriever(test_qa_list, sparse_ret_for_eval, da_idf_recall)\n",
    "\n",
    "df = pd.DataFrame({\n",
    "    'IDF-Recall': [sparse_score['idf_recall'], dense_score['idf_recall']],\n",
    "    'Time/Query': [sparse_score['retrieval_time'], dense_score['retrieval_time']],\n",
    "}, index=[ 'Sparse', 'Dense'])\n",
    "df.style.format({'IDF-Recall': \"{:.4f}\", 'Time/Query': \"{:.4f}\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This metric obviously has limitations. Since it only looks at individual word matches, it is oblivious to *word order* or *paraphrases* among others. However, we can be encouraged by the fact that the dense retriever not only yields **higher IDF-recall**, it also takes **less than a third of the time** of the ElasticSearch-based system! Considering these results, we can confidently use it for the next part: training the sequence-to-sequence answer generation system."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='generation'></a>\n",
    "# 5. Generating Answers with a Sequence-to-Sequence Model\n",
    "\n",
    "Now that we know how to create an evidence document with supporting information for a given question, let's look into training the second component of our system: the **answer generation module**. We will instantiate it as a sequence-to-sequence model which uses the [BART](https://arxiv.org/abs/1910.13461) architecture, and initialize it with the [bart-large pretrained weights](https://huggingface.co/facebook/bart-large).  \n",
    "\n",
    "In short, the [BART paper](https://arxiv.org/abs/1910.13461) uses a denoising auto-encoder style objective to pre-train an encoder-decoder model (similarly to how masked language modeling is used to pre-trained BERT-style encoders). Among other applications, they show that large-scale pre-training with their objective followed by fine-tuning on ELI5 data yields the state-of-the-art ROUGE performance for the original version of the dataset (which uses pre-computed support documents made from CommonCrawl pages).\n",
    "\n",
    "We provide the concatenation of the question and support document as input to the model, and train the decoder to minimize the perplexity of the gold answer. One notable choice is that **we train the model using all high-scoring answers in the training set**, so the model will see several instances of the same question-document input with different outputs. The supporting passages are separated by a special token `<P>`, so the input for our running example will look like:\n",
    "\n",
    "> question: Why does water heated to room temperature feel colder than the air around it? context: \\\\<P\\> when the skin is completely wet. The body continuously loses ... this heat comes from the liquid itself and the surrounding gas and surfaces. \\\\<P\\> protected by a glass panel. Consequently, these types of collectors... Since heat loss due to convection cannot cross a vacuum, it forms an efficient isolation mechanism to keep heat inside the collector pipes. Since two flat \\\\<P\\> ... \\\\<P\\> changes. Conduction On... Fluids—especially gases—are less conductive. Thermal contact conductance is the study of heat conduction between solid bodies in contact. The process of heat transfer\n",
    "\n",
    "The first thing we do is pre-compute the support documents for the training and validation sets so we can use all available GPUs to train the sequence-to-sequence model. The model is then trained with the `train_qa_s2s` function in `eli5_utils.py`. A 16GB GPU accomodates up to two examples at a time, so here is the code to train the model using 4 GPUs with `torch.nn.DataPArallel`. One epoch should take about 18 hours:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pre-computing support documents\n",
    "eli5_train_docs = []\n",
    "for example in eli5['train_eli5']:\n",
    "    support_doc, dense_res_list = query_qa_dense_index(\n",
    "        example['title'], qar_model, qar_tokenizer, wiki40b_snippets, wiki40b_gpu_index, n_results=n_ret\n",
    "    )\n",
    "    eli5_train_docs += [(example['q_id'], support_doc, dense_res_list)]\n",
    "\n",
    "eli5_valid_docs = []\n",
    "for example in eli5['validation_eli5']:\n",
    "    support_doc, dense_res_list = query_qa_dense_index(\n",
    "        example['title'], qar_model, qar_tokenizer, wiki40b_snippets, wiki40b_gpu_index, n_results=n_ret\n",
    "    )\n",
    "    eli5_valid_docs += [(example['q_id'], support_doc, dense_res_list)]\n",
    "\n",
    "# training loop proper\n",
    "class ArgumentsS2S():\n",
    "    def __init__(self):\n",
    "        self.batch_size = 8\n",
    "        self.backward_freq = 16\n",
    "        self.max_length = 1024\n",
    "        self.print_freq = 100\n",
    "        self.model_save_name = \"seq2seq_models/eli5_bart_model\"\n",
    "        self.learning_rate = 2e-4\n",
    "        self.num_epochs = 3\n",
    "\n",
    "s2s_args = ArgumentsS2S()\n",
    "\n",
    "eli5_train_docs = json.load(open('precomputed/eli5_train_precomputed_dense_docs.json'))\n",
    "eli5_valid_docs = json.load(open('precomputed/eli5_valid_precomputed_dense_docs.json'))\n",
    "s2s_train_dset = ELI5DatasetS2S(eli5['train_eli5'], document_cache=dict([(k, d) for k, d, src_ls in eli5_train_docs]))\n",
    "s2s_valid_dset = ELI5DatasetS2S(eli5['validation_eli5'], document_cache=dict([(k, d) for k, d, src_ls in eli5_valid_docs]), training=False)\n",
    "\n",
    "qa_s2s_tokenizer, pre_model = make_qa_s2s_model(\n",
    "    model_name=\"facebook/bart-large\",\n",
    "    from_file=None,\n",
    "    device=\"cuda:0\"\n",
    ")\n",
    "qa_s2s_model = torch.nn.DataParallel(pre_model)\n",
    "\n",
    "train_qa_s2s(qa_s2s_model, qa_s2s_tokenizer, s2s_train_dset, s2s_valid_dset, s2s_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, if you don't want to train the model yourself, we made trained weights available on the [Hugging Face model repository](https://huggingface.co/models) , which you can download with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_s2s_tokenizer = AutoTokenizer.from_pretrained('yjernite/bart_eli5')\n",
    "qa_s2s_model = AutoModelForSeq2SeqLM.from_pretrained('yjernite/bart_eli5').to('cuda:0')\n",
    "_ = qa_s2s_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have everything we need to answer any question! Now let's try the full system on our running example along with the first four questions of the test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row0_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row0_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row1_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row1_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row2_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row2_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row3_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row3_col1 {\n",
       "            text-align:  left;\n",
       "        }    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row4_col0 {\n",
       "            text-align:  left;\n",
       "        }    #T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row4_col1 {\n",
       "            text-align:  left;\n",
       "        }</style><table id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >Question</th>        <th class=\"col_heading level0 col1\" >Answer</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row0_col0\" class=\"data row0 col0\" >Why does water heated to room temperature feel colder than the air around it?</td>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row0_col1\" class=\"data row0 col1\" >It doesn't feel colder than the air around it, it feels colder than your body temperature. Water is a better conductor of heat than air, so it takes more energy to heat it up than air is to cool it down. So when you heat water to room temperature, it takes away more energy from your body than air does.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076level0_row1\" class=\"row_heading level0 row1\" >1</th>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row1_col0\" class=\"data row1 col0\" >Why do you get chills/goosebumps from hearing large crowds sing along to songs?</td>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row1_col1\" class=\"data row1 col1\" >I get goosebumps when I sing along to a song. I don't know what causes it, but it happens to me all the time. I think it has something to do with the fact that when you sing along with a song, your brain releases endorphins, which make you feel good.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076level0_row2\" class=\"row_heading level0 row2\" >2</th>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row2_col0\" class=\"data row2 col0\" >How did studded leather and heavy eye makeup come to be the Hollywood dress code for dystopian, post-apocalyptic societies?</td>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row2_col1\" class=\"data row2 col1\" >Studded leather and heavy eye makeup have been around for a long time. It's not a new thing, it's been a Hollywood thing since the 80s and 90s.  URL_0   >  Studded and eye makeup were popularized by Alfred Hitchcock in the 1950s and 1960s, and a number of occurrences of the style in films were mentioned - Charlize Theron in Aeon Flux, Milla Jovovich in the fourth and fifth film of the Fourth and fifth movie of the Star Wars franchise, Angelina Jolie in The Last Airbender, and Raquel Croft in Raiders of the Lost Ark. According to Alfred Hitchcock's To Catch a Thief, there was a trend of women wearing pouty lips, Pouty Lips, pursed lips, etc. and a lot of depictions of the female role of the protagonist in a dystopian, post-apocalyptic society.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076level0_row3\" class=\"row_heading level0 row3\" >3</th>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row3_col0\" class=\"data row3 col0\" >What's the difference between a bush, a shrub, and a tree?</td>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row3_col1\" class=\"data row3 col1\" >A tree is a living thing. A shrub is a kind of plant. A bush is a type of plant that grows in the ground. A tree has a trunk, a branch, and a trunk. The trunk is the main part of the tree, the branch is the part that grows into the trunk.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076level0_row4\" class=\"row_heading level0 row4\" >4</th>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row4_col0\" class=\"data row4 col0\" >Why is it hard to breathe with a strong air gust blowing straight at your face?</td>\n",
       "                        <td id=\"T_a4ac9ec6_acec_11ea_ba32_1b67035d3076row4_col1\" class=\"data row4 col1\" >It's not hard to breathe with a strong air gust blowing straight at your face. It's hard to breath when the wind is blowing in the opposite direction. The wind is pushing the air away from your face, so it's harder for your lungs to get the air they need to work against the wind.</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7fd98edb9910>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "questions = []\n",
    "answers = []\n",
    "\n",
    "for i in [12345] + [j for j in range(4)]:\n",
    "    # create support document with the dense index\n",
    "    question = eli5['test_eli5'][i]['title']\n",
    "    doc, res_list = query_qa_dense_index(\n",
    "        question, qar_model, qar_tokenizer,\n",
    "        wiki40b_snippets, wiki40b_gpu_index, device='cuda:1'\n",
    "    )\n",
    "    # concatenate question and support document into BART input\n",
    "    question_doc = \"question: {} context: {}\".format(question, doc)\n",
    "    # generate an answer with beam search\n",
    "    answer = qa_s2s_generate(\n",
    "            question_doc, qa_s2s_model, qa_s2s_tokenizer,\n",
    "            num_answers=1,\n",
    "            num_beams=8,\n",
    "            min_len=64,\n",
    "            max_len=256,\n",
    "            max_input_length=1024,\n",
    "            device=\"cuda:0\"\n",
    "    )[0]\n",
    "    questions += [question]\n",
    "    answers += [answer]\n",
    "\n",
    "df = pd.DataFrame({\n",
    "    'Question': questions,\n",
    "    'Answer': answers,\n",
    "})\n",
    "df.style.set_properties(**{'text-align': 'left'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/fireworks.gif\" width=\"1000\" align=\"center\"/>  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We made it, and a lot of these answers actually make sense! The model seems to sometimes struggle with coherence and with starting some of the answers, but we're getting some pretty good information overall.\n",
    "\n",
    "The last thing we'll do is see how we can get a quantitative evaluation of the model performance. Here, we'll use the ROUGE implementation provided in the `nlp` library.  \n",
    "\n",
    "Note that it is a different implementation than the one used in the [BART](https://arxiv.org/abs/1910.13461) and [ELI5](https://arxiv.org/abs/1907.09190) papers: the [rouge](https://pypi.org/project/rouge/) Python package they use normalises all numerical values, among other pre-processing choices, leading to higher numbers. We reproduce their evaluation in the Appendix section, but recommend using the more sensitive metric provided by the `nlp` package, which can be computed with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted = []\n",
    "reference = []\n",
    "\n",
    "# Generate answers for the full test set\n",
    "for i in range(eli5['test_eli5'].num_rows):\n",
    "    # create support document with the dense index\n",
    "    question = eli5['test_eli5'][i]['title']\n",
    "    doc, res_list = query_qa_dense_index(\n",
    "        question, qar_model, qar_tokenizer,\n",
    "        wiki40b_snippets, wiki40b_gpu_index, device='cuda:1'\n",
    "    )\n",
    "    # concatenate question and support document into BART input\n",
    "    question_doc = \"question: {} context: {}\".format(question, doc)\n",
    "    # generate an answer with beam search\n",
    "    answer = qa_s2s_generate(\n",
    "            question_doc, qa_s2s_model, qa_s2s_tokenizer,\n",
    "            num_answers=1,\n",
    "            num_beams=8,\n",
    "            min_len=96,\n",
    "            max_len=256,\n",
    "            max_input_length=1024,\n",
    "            device=\"cuda:0\"\n",
    "    )[0]\n",
    "    predicted += [answer]\n",
    "    reference += [eli5['test_eli5'][i]['answers']['text'][0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "</style><table id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >rouge1</th>        <th class=\"col_heading level0 col1\" >rouge2</th>        <th class=\"col_heading level0 col2\" >rougeL</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806level0_row0\" class=\"row_heading level0 row0\" >P</th>\n",
       "                        <td id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806row0_col0\" class=\"data row0 col0\" >0.3025</td>\n",
       "                        <td id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806row0_col1\" class=\"data row0 col1\" >0.0609</td>\n",
       "                        <td id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806row0_col2\" class=\"data row0 col2\" >0.1708</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806level0_row1\" class=\"row_heading level0 row1\" >R</th>\n",
       "                        <td id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806row1_col0\" class=\"data row1 col0\" >0.2946</td>\n",
       "                        <td id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806row1_col1\" class=\"data row1 col1\" >0.0587</td>\n",
       "                        <td id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806row1_col2\" class=\"data row1 col2\" >0.1797</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806level0_row2\" class=\"row_heading level0 row2\" >F</th>\n",
       "                        <td id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806row2_col0\" class=\"data row2 col0\" >0.2561</td>\n",
       "                        <td id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806row2_col1\" class=\"data row2 col1\" >0.0504</td>\n",
       "                        <td id=\"T_7850660e_acb1_11ea_bac0_cb508f4c6806row2_col2\" class=\"data row2 col2\" >0.1489</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f098709df10>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Compare each generation to the fist answer from the dataset\n",
    "nlp_rouge = nlp.load_metric('rouge')\n",
    "\n",
    "scores = nlp_rouge.compute(\n",
    "    predicted, reference,\n",
    "    rouge_types=['rouge1', 'rouge2', 'rougeL', 'rougeLsum'],\n",
    "    use_agregator=True, use_stemmer=False\n",
    ")\n",
    "df = pd.DataFrame({\n",
    "    'rouge1': [scores['rouge1'].mid.precision, scores['rouge1'].mid.recall, scores['rouge1'].mid.fmeasure],\n",
    "    'rouge2': [scores['rouge2'].mid.precision, scores['rouge2'].mid.recall, scores['rouge2'].mid.fmeasure],\n",
    "    'rougeL': [scores['rougeL'].mid.precision, scores['rougeL'].mid.recall, scores['rougeL'].mid.fmeasure],\n",
    "}, index=[ 'P', 'R', 'F'])\n",
    "df.style.format({'rouge1': \"{:.4f}\", 'rouge2': \"{:.4f}\", 'rougeL': \"{:.4f}\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it for today! And once again, if you want to play with the model a bit more and ask it whatever question comes to mind, please feel free to head over to:\n",
    "# [**Our Live Demo!**](https://huggingface.co/qa/)  \n",
    "\n",
    "Thank you for reading!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Appendix:\n",
    "\n",
    "Here we reproduce the ROUGE evaluation from the original [ELI5 paper](https://arxiv.org/abs/1907.09190) to be able to comparable our performance to theirs. Our generation setting leads to lower ROUGE-1 and ROUGE-2 than the state-of-the-art reported in [BART](https://arxiv.org/abs/1910.13461) (30.6 and 6.2 respectively), and higher ROUGE-L (24.3)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "</style><table id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >rouge1</th>        <th class=\"col_heading level0 col1\" >rouge2</th>        <th class=\"col_heading level0 col2\" >rougeL</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806level0_row0\" class=\"row_heading level0 row0\" >P</th>\n",
       "                        <td id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806row0_col0\" class=\"data row0 col0\" >0.3254</td>\n",
       "                        <td id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806row0_col1\" class=\"data row0 col1\" >0.0680</td>\n",
       "                        <td id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806row0_col2\" class=\"data row0 col2\" >0.3251</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806level0_row1\" class=\"row_heading level0 row1\" >R</th>\n",
       "                        <td id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806row1_col0\" class=\"data row1 col0\" >0.3118</td>\n",
       "                        <td id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806row1_col1\" class=\"data row1 col1\" >0.0631</td>\n",
       "                        <td id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806row1_col2\" class=\"data row1 col2\" >0.2560</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806level0_row2\" class=\"row_heading level0 row2\" >F</th>\n",
       "                        <td id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806row2_col0\" class=\"data row2 col0\" >0.2729</td>\n",
       "                        <td id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806row2_col1\" class=\"data row2 col1\" >0.0551</td>\n",
       "                        <td id=\"T_59d1df5e_acb2_11ea_bac0_cb508f4c6806row2_col2\" class=\"data row2 col2\" >0.2583</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f09776cc310>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from nltk import PorterStemmer\n",
    "from rouge import Rouge\n",
    "from spacy.lang.en import English\n",
    "from time import time\n",
    "\n",
    "stemmer = PorterStemmer()\n",
    "rouge = Rouge()\n",
    "tokenizer = English().Defaults.create_tokenizer()\n",
    "\n",
    "def compute_rouge_eli5(compare_list):\n",
    "    preds = [\" \".join([stemmer.stem(str(w))\n",
    "                       for w in tokenizer(pred)])\n",
    "             for gold, pred in compare_list]\n",
    "    golds = [\" \".join([stemmer.stem(str(w))\n",
    "                       for w in tokenizer(gold)])\n",
    "             for gold, pred in compare_list]\n",
    "    scores = rouge.get_scores(preds, golds, avg=True)\n",
    "    return scores\n",
    "\n",
    "\n",
    "compare_list = [(g, p) for p, g in zip(predicted, reference)]\n",
    "scores = compute_rouge_eli5(compare_list)\n",
    "df = pd.DataFrame({\n",
    "    'rouge1': [scores['rouge-1']['p'], scores['rouge-1']['r'], scores['rouge-1']['f']],\n",
    "    'rouge2': [scores['rouge-2']['p'], scores['rouge-2']['r'], scores['rouge-2']['f']],\n",
    "    'rougeL': [scores['rouge-l']['p'], scores['rouge-l']['r'], scores['rouge-l']['f']],\n",
    "}, index=[ 'P', 'R', 'F'])\n",
    "df.style.format({'rouge1': \"{:.4f}\", 'rouge2': \"{:.4f}\", 'rougeL': \"{:.4f}\"})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
